package com.study.qdrant.service.impl;

import com.study.qdrant.service.QDrantService;
import io.metaloom.qdrant.client.grpc.QDrantGRPCClient;
import io.metaloom.qdrant.client.grpc.proto.Collections;
import io.metaloom.qdrant.client.grpc.proto.JsonWithInt;
import io.metaloom.qdrant.client.grpc.proto.Points;
import io.metaloom.qdrant.client.util.ModelHelper;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

import static io.metaloom.qdrant.client.util.ModelHelper.value;

public class QDrantGrpcServiceImpl extends QDrantBaseService implements QDrantService {

    public void grpc() throws ExecutionException, InterruptedException {
        int port = 6334;//qdrant.grpcPort(); // Default: 6334
        try (QDrantGRPCClient client = QDrantGRPCClient.builder()
                .setHostname(host)
                .setPort(port)
                .build()) {

            // Define the collection to store vectors
            Collections.VectorParams params = Collections.VectorParams.newBuilder()
                    .setSize(4)
                    .setDistance(Collections.Distance.Euclid)
                    .build();

            // Add the params to a map
            Collections.VectorParamsMap paramsMap = Collections.VectorParamsMap.newBuilder()
                    .putMap("firstVector", params)
                    .putMap("secondVector", params)
                    .build();
            // Create new collections - blocking
            client.createCollection("test1", paramsMap)
                    .sync();
            // .. or via Future API
//            client.createCollection("test2", params).async().get();
//            // .. or via RxJava API
//            client.createCollection("test3", params).rx().blockingGet();

            // Insert a new vectors
            for (int i = 0; i < 10; i++) {
                // Vector of the point
                float[] vector = new float[]{0.43f + i, 0.1f, 0.61f, 1.45f - i};

                // Payload of the point
                Map<String, JsonWithInt.Value> payload = new HashMap<>();
                payload.put("color", ModelHelper.value("blue"));

                // Now construct the point
                Points.PointStruct point = ModelHelper.namedPoint(42L + i, "firstVector", vector, payload);
                // .. and insert it
                client.upsertPoint("test1", point, true).sync();
            }

            // Count points
            long nPoints = client.countPoints("test1", null, true).sync().getResult().getCount();

            // Now run KNN search
            float[] searchVector = new float[]{0.43f, 0.09f, 0.41f, 1.35f};
            List<Points.ScoredPoint> searchResults = client.searchPoints("test1", "firstVector", searchVector, 2, null)
                    .sync()
                    .getResultList();
            for (Points.ScoredPoint result : searchResults) {
                System.out.println("Found: [" + result.getId().getNum() + "] " + result.getScore());
            }

            // Invoke backup via Snapshot API
            client.createSnapshot("test1").sync();
        }
    }


    private Points.Vectors createVector(int size) {
        return Points.Vectors.newBuilder().build();
    }

    public void process() {
        int grpcPort = 6334;
        try (QDrantGRPCClient client = QDrantGRPCClient.builder().setHostname(host).setPort(grpcPort).build()) {

            final String collectionName = "test_collection2";
            //  创建集合: 将把所有矢量数据存储在一个 Qdrant 集合中。我们称之为.此集合将使用点积距离度量来比较向量。
            Collections.GetCollectionInfoResponse res = client.loadCollections(collectionName).sync();
            // TODO 判断是否存在
            if(res.isInitialized()) {
                client.deleteCollection(collectionName,null).sync();
            }
            // 随机生成向量:个数、距离类型
            Collections.VectorParams params = Collections.VectorParams.newBuilder()
                    .setDistance(Collections.Distance.Euclid)
                    .setSize(4)
                    .build();
            Collections.VectorParamsMap paramsMap = Collections.VectorParamsMap.newBuilder()
                    .putMap("firstVector", params)
                    .putMap("secondVector", params)
                    .build();
            Collections.CollectionOperationResponse response = client.createCollection(collectionName, paramsMap).sync();
            System.out.println(response);

            // 添加

            //Points.Vectors.newBuilder().build()
            float[] vector = new float[]{0.05f, 0.61f, 0.76f, 0.74f};
            // Payload  Map.of("city", ModelHelper.value("Berlin"))
            Map<String, JsonWithInt.Value> payload = new HashMap<>();
            payload.put("color", ModelHelper.value("blue"));

            Points.PointStruct p1 = ModelHelper.namedPoint(42L, "firstVector", vector, payload);
            Points.PointsOperationResponse operationInfo = client.upsertPoint(collectionName, p1, true)
                    .sync();
            System.out.println(operationInfo);

            // 查询 KNN
            float[] searchVector = new float[]{0.43f, 0.09f, 0.41f, 1.35f};
            List<Points.ScoredPoint> searchResults = client.searchPoints(collectionName, "firstVector", searchVector, 10, null)
                    .sync()
                    .getResultList();
            for (Points.ScoredPoint result : searchResults) {
                System.out.println("Found: [" + result.getId().getNum() + "] " + result.getScore());
            }

            // Count points
            long nPoints = client.countPoints(collectionName, null, true).sync().getResult().getCount();
            System.out.println(nPoints);
        }
    }

    public static void main(String[] args) throws Exception {
        QDrantGrpcServiceImpl test = new QDrantGrpcServiceImpl();
//        test.grpc();
        test.process();
    }
}
